package cert

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"network-firewall-go/internal/types"
)

// Manager 证书管理器
type Manager struct {
	rootCA     *x509.Certificate
	rootKey    *rsa.PrivateKey
	certCache  map[string]*tls.Certificate
	cacheMutex sync.RWMutex
	certDir    string
}

// NewManager 创建新的证书管理器
func NewManager(certDir string) (*Manager, error) {
	if err := os.MkdirAll(certDir, 0755); err != nil {
		return nil, fmt.Errorf("创建证书目录失败: %w", err)
	}

	manager := &Manager{
		certCache: make(map[string]*tls.Certificate),
		certDir:   certDir,
	}

	// 初始化根证书
	if err := manager.initRootCA(); err != nil {
		return nil, fmt.Errorf("初始化根证书失败: %w", err)
	}

	return manager, nil
}

// initRootCA 初始化根证书
func (m *Manager) initRootCA() error {
	rootCertPath := filepath.Join(m.certDir, "root-ca.crt")
	rootKeyPath := filepath.Join(m.certDir, "root-ca.key")

	// 检查是否已存在根证书
	if _, err := os.Stat(rootCertPath); err == nil {
		if _, err := os.Stat(rootKeyPath); err == nil {
			return m.loadRootCA(rootCertPath, rootKeyPath)
		}
	}

	// 生成新的根证书
	return m.generateRootCA(rootCertPath, rootKeyPath)
}

// loadRootCA 加载现有的根证书
func (m *Manager) loadRootCA(certPath, keyPath string) error {
	// 读取证书文件
	certPEM, err := os.ReadFile(certPath)
	if err != nil {
		return fmt.Errorf("读取根证书文件失败: %w", err)
	}

	certBlock, _ := pem.Decode(certPEM)
	if certBlock == nil {
		return fmt.Errorf("解析根证书PEM失败")
	}

	cert, err := x509.ParseCertificate(certBlock.Bytes)
	if err != nil {
		return fmt.Errorf("解析根证书失败: %w", err)
	}

	// 读取私钥文件
	keyPEM, err := os.ReadFile(keyPath)
	if err != nil {
		return fmt.Errorf("读取根证书私钥文件失败: %w", err)
	}

	keyBlock, _ := pem.Decode(keyPEM)
	if keyBlock == nil {
		return fmt.Errorf("解析根证书私钥PEM失败")
	}

	key, err := x509.ParsePKCS1PrivateKey(keyBlock.Bytes)
	if err != nil {
		return fmt.Errorf("解析根证书私钥失败: %w", err)
	}

	m.rootCA = cert
	m.rootKey = key

	fmt.Printf("✅ 已加载现有根证书: %s\n", cert.Subject.CommonName)
	return nil
}

// generateRootCA 生成新的根证书
func (m *Manager) generateRootCA(certPath, keyPath string) error {
	fmt.Println("🔐 生成新的根CA证书...")

	// 生成私钥
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return fmt.Errorf("生成根证书私钥失败: %w", err)
	}

	// 创建证书模板
	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName:         "Network Firewall Proxy Root CA",
			Country:            []string{"CN"},
			Province:           []string{"Beijing"},
			Locality:           []string{"Beijing"},
			Organization:       []string{"Network Firewall Proxy"},
			OrganizationalUnit: []string{"IT Department"},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(10 * 365 * 24 * time.Hour), // 10年
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA:                  true,
		MaxPathLen:            0,
	}

	// 自签名
	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
	if err != nil {
		return fmt.Errorf("创建根证书失败: %w", err)
	}

	// 解析生成的证书
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return fmt.Errorf("解析生成的根证书失败: %w", err)
	}

	// 保存证书文件
	certFile, err := os.Create(certPath)
	if err != nil {
		return fmt.Errorf("创建根证书文件失败: %w", err)
	}
	defer certFile.Close()

	if err := pem.Encode(certFile, &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certDER,
	}); err != nil {
		return fmt.Errorf("写入根证书文件失败: %w", err)
	}

	// 保存私钥文件
	keyFile, err := os.Create(keyPath)
	if err != nil {
		return fmt.Errorf("创建根证书私钥文件失败: %w", err)
	}
	defer keyFile.Close()

	if err := pem.Encode(keyFile, &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}); err != nil {
		return fmt.Errorf("写入根证书私钥文件失败: %w", err)
	}

	m.rootCA = cert
	m.rootKey = key

	fmt.Printf("✅ 根CA证书生成完成\n")
	fmt.Printf("📁 证书文件: %s\n", certPath)
	fmt.Printf("🔑 私钥文件: %s\n", keyPath)
	fmt.Printf("⚠️  请将根证书安装到系统信任存储中以启用HTTPS拦截\n")

	return nil
}

// GetCertificateForHost 为指定主机获取证书
func (m *Manager) GetCertificateForHost(hostname string) (*tls.Certificate, error) {
	m.cacheMutex.RLock()
	if cert, exists := m.certCache[hostname]; exists {
		m.cacheMutex.RUnlock()
		return cert, nil
	}
	m.cacheMutex.RUnlock()

	// 生成新证书
	cert, err := m.generateCertificateForHost(hostname)
	if err != nil {
		return nil, err
	}

	// 缓存证书
	m.cacheMutex.Lock()
	m.certCache[hostname] = cert
	m.cacheMutex.Unlock()

	return cert, nil
}

// generateCertificateForHost 为指定主机生成证书
func (m *Manager) generateCertificateForHost(hostname string) (*tls.Certificate, error) {
	fmt.Printf("🔐 为 %s 生成证书...\n", hostname)

	// 生成私钥
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("生成证书私钥失败: %w", err)
	}

	// 创建证书模板
	template := x509.Certificate{
		SerialNumber: big.NewInt(time.Now().UnixNano()),
		Subject: pkix.Name{
			CommonName:         hostname,
			Country:            []string{"CN"},
			Province:           []string{"Beijing"},
			Locality:           []string{"Beijing"},
			Organization:       []string{"Network Firewall Proxy"},
			OrganizationalUnit: []string{"IT Department"},
		},
		NotBefore:   time.Now(),
		NotAfter:    time.Now().Add(365 * 24 * time.Hour), // 1年
		KeyUsage:    x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSNames:    []string{hostname},
	}

	// 如果是IP地址，添加到IPAddresses
	if ip := net.ParseIP(hostname); ip != nil {
		template.IPAddresses = []net.IP{ip}
	} else {
		// 添加通配符域名支持
		if hostname != "localhost" && hostname != "127.0.0.1" {
			template.DNSNames = append(template.DNSNames, "*."+hostname)
		}
	}

	// 使用根CA签名
	certDER, err := x509.CreateCertificate(rand.Reader, &template, m.rootCA, &key.PublicKey, m.rootKey)
	if err != nil {
		return nil, fmt.Errorf("创建证书失败: %w", err)
	}

	// 创建TLS证书
	cert := &tls.Certificate{
		Certificate: [][]byte{certDER},
		PrivateKey:  key,
	}

	fmt.Printf("✅ %s 证书生成完成\n", hostname)
	return cert, nil
}

// GetSNICallback 获取SNI回调函数
func (m *Manager) GetSNICallback() func(*tls.ClientHelloInfo) (*tls.Certificate, error) {
	return func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
		hostname := hello.ServerName
		if hostname == "" {
			hostname = "0.0.0.0"
		}

		return m.GetCertificateForHost(hostname)
	}
}

// GetRootCACertificate 获取根CA证书信息
func (m *Manager) GetRootCACertificate() *types.CertificateInfo {
	if m.rootCA == nil {
		return nil
	}

	return &types.CertificateInfo{
		Subject:     m.rootCA.Subject.String(),
		Issuer:      m.rootCA.Issuer.String(),
		NotBefore:   m.rootCA.NotBefore,
		NotAfter:    m.rootCA.NotAfter,
		DNSNames:    m.rootCA.DNSNames,
		Fingerprint: fmt.Sprintf("%x", m.rootCA.Raw),
	}
}

// GetRootCAPEM 获取根CA证书的PEM格式
func (m *Manager) GetRootCAPEM() ([]byte, error) {
	if m.rootCA == nil {
		return nil, fmt.Errorf("根CA证书未初始化")
	}

	return pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: m.rootCA.Raw,
	}), nil
}

// ClearCache 清理证书缓存
func (m *Manager) ClearCache() {
	m.cacheMutex.Lock()
	defer m.cacheMutex.Unlock()

	m.certCache = make(map[string]*tls.Certificate)
	fmt.Println("🧹 证书缓存已清理")
}

// GetCacheStats 获取缓存统计信息
func (m *Manager) GetCacheStats() map[string]interface{} {
	m.cacheMutex.RLock()
	defer m.cacheMutex.RUnlock()

	return map[string]interface{}{
		"cached_certificates": len(m.certCache),
		"has_root_ca":         m.rootCA != nil,
		"root_ca_subject": func() string {
			if m.rootCA != nil {
				return m.rootCA.Subject.CommonName
			}
			return ""
		}(),
	}
}
